## 模板1
import os


def cuda_builder(target, source, env):
    # nvcc 编译 .cu 文件
    # 1.编译
    objs = []
    for file in source:
        file = str(file)
        if file.endswith(".cu"):
            file_name, _ = os.path.splitext(file)
            file_obj = f"{file_name}.o"
            build_cmd = f"nvcc -c {file} -o {file_obj} {env['CCFLAGS']}"
            ret = env.Execute(build_cmd)
            if ret != 0:
                raise Exception(f"Failed to compile {file}")
            objs.append(file_obj)

    # 2.链接
    target_file = str(target[0])
    link_cmd = f"nvcc -o {target_file} {' '.join(objs)}"
    ret = env.Execute(link_cmd)
    if ret != 0:
        raise Exception(f"Failed to link {target_file} {file}")


# 创建自定义构建器
cuda_builder = Builder(action=cuda_builder, suffix=".out", src_suffix=".cu")

env = Environment()
# env["CCFLAGS"] = "-g"                 # gdb 调试开关
env.Append(BUILDERS={"cudaBuilder": cuda_builder})
env.cudaBuilder("1_pinned_mem也内存锁定.cu")
